
from ctypes import *
import base64
import hashlib
from io import BytesIO
from Cryptodome.Random import get_random_bytes
from Cryptodome.Cipher import AES

class LibuECCBuffer32(Structure):
	_fields_ = [("buffer", c_ubyte * 32)]
	def fromBytes(self, bytes):
		self.buffer = (c_ubyte * 32).from_buffer_copy(bytes)
	def toBase64(self):
		return base64.encodebytes(bytearray(self.buffer))
	def fromBase64(self,b64_str):
		self.fromBytes(base64.decodebytes(b64_str))
	def toBytes(self):
		return bytes(bytearray(self.buffer))

class LibuECCBuffer64(Structure):
	_fields_ = [("buffer", c_ubyte * 64)]
	def fromBytes(self, bytes):
		self.buffer = (c_ubyte * 64).from_buffer_copy(bytes)
	def toBase64(self):
		return base64.encodebytes(bytearray(self.buffer))
	def fromBase64(self,b64_str):
		self.fromBytes(base64.decodebytes(b64_str))
	def toBytes(self):
		return bytes(bytearray(self.buffer))

class LibuECC(object):
	__LibuECCPrivateKeyPointer = POINTER(LibuECCBuffer32)
	__LibuECCPublicKeyPointer = POINTER(LibuECCBuffer64)
	__LibuECCSha256Pointer = POINTER(LibuECCBuffer32)
	__LibuECCSignedData = POINTER(LibuECCBuffer64)
	__LibuECCSharedSecret = POINTER(LibuECCBuffer32)
	__uECCdll = WinDLL("LibuECC.dll")
	__Dll_return_success = 0

	__uECC_get_verion = __uECCdll['get_version']

	__uECC_generate_key_pair = __uECCdll['LibuECC_generate_key_pair']
	__uECC_generate_key_pair.argtypes = [__LibuECCPrivateKeyPointer, __LibuECCPublicKeyPointer]

	__uECC_get_public_key = __uECCdll['LibuECC_get_publickey']
	__uECC_get_public_key.argtypes = [__LibuECCPrivateKeyPointer,
		__LibuECCPublicKeyPointer]

	__uECC_get_shared_secret = __uECCdll['LibuECC_get_shared_secret']
	__uECC_get_shared_secret.argtypes = [__LibuECCPrivateKeyPointer,
		__LibuECCPublicKeyPointer,
		__LibuECCSharedSecret]

	__uECC_is_on_Curve = __uECCdll['LibuECC_is_on_Curve']
	__uECC_is_on_Curve.argtypes = [__LibuECCPublicKeyPointer]

	__uECC_sign = __uECCdll['LibuECC_Sign']
	__uECC_sign.argtypes = [__LibuECCPrivateKeyPointer,
		__LibuECCSha256Pointer,
		__LibuECCSignedData]

	__uECC_verify = __uECCdll['LibuECC_verify']
	__uECC_verify.argtypes = [__LibuECCPublicKeyPointer,
		__LibuECCSha256Pointer,
		__LibuECCSignedData,
		POINTER(c_bool)]

	def __init__(self):
		self.__private_key = LibuECCBuffer32()
		self.__public_key = LibuECCBuffer64()
		self.__shared_secret = LibuECCBuffer32()

	@staticmethod
	def getVersion():
		return LibuECC.__uECC_get_verion()

	@staticmethod
	def __getStrSha256Buffer(target_str):
		sh = hashlib.sha256()
		sh.update(target_str.encode('utf-8'))
		buffer = LibuECCBuffer32()
		buffer.fromBytes(sh.digest())
		return buffer

	def writePrivateKeyBase64(self, pri_base64):
		self.__private_key.fromBase64(pri_base64)

	def __readPrivateKeyBase64(self):
		return self.__private_key.toBase64()

	def writePublicKeyBase64(self, pub_base64):
		self.__public_key.fromBase64(pub_base64)

	def readPublicKeyBase64(self):
		return self.__public_key.toBase64()

	def readSharedSecretBytes(self):
		return self.__shared_secret.toBytes();

	def reGeneratePublicKey(self):
		ret = LibuECC.__uECC_get_public_key(self.__private_key, self.__public_key)
		if ret != LibuECC.__Dll_return_success:
			raise BufferError

	def generateKeyPair(self):
		ret = LibuECC.__uECC_generate_key_pair(self.__private_key, self.__public_key)
		if ret != LibuECC.__Dll_return_success:
			raise BufferError

	def Sign(self, str_sign):
		sha = LibuECC.__getStrSha256Buffer(str_sign)
		signed_data = LibuECCBuffer64();
		ret = LibuECC.__uECC_sign(self.__private_key, sha, signed_data)
		if ret != LibuECC.__Dll_return_success:
			raise BufferError
		return signed_data.toBase64()

	def Verify(self, str_toCheck, str_signed_data):
		sha = LibuECC.__getStrSha256Buffer(str_toCheck)
		signed_data = LibuECCBuffer64()
		signed_data.fromBase64(str_signed_data);
		isSignChecked = c_bool()
		ret = LibuECC.__uECC_verify(self.__public_key,sha,signed_data,byref(isSignChecked))
		if isSignChecked:
			return True
		else:
			return False

	def SharedSecret(self, str_remote_public_key_base64):
		remote_pub_key = LibuECCBuffer64()
		remote_pub_key.fromBase64(str_remote_public_key_base64)
		ret = LibuECC.__uECC_get_shared_secret(self.__private_key, remote_pub_key, self.__shared_secret)
		if ret != LibuECC.__Dll_return_success:
			raise BufferError
		return self.__shared_secret.toBase64()

	def Encrypt(self, msg, nonce_len = AES.block_size):
		iv = get_random_bytes(nonce_len)
		cipher = AES.new(self.__shared_secret.toBytes(), AES.MODE_GCM, iv)
		encryped = cipher.encrypt(msg.encode('utf-8'))
		return base64.b64encode(iv + encryped)

	def Decrype(self, crypted, nonce_len = AES.block_size):
		raw = BytesIO(base64.b64decode(crypted))
		iv = raw.read(nonce_len)
		encrypted = raw.read()
		cipher = AES.new(self.__shared_secret.toBytes(), AES.MODE_GCM, iv)
		decryped = cipher.decrypt(encrypted)
		return decryped.decode('utf-8')
		

if __name__ == "__main__":
	uSign = LibuECC()
	print("getVersion:" + str(uSign.getVersion()))
	uSign.generateKeyPair()
	str_base64_pub_check = uSign.readPublicKeyBase64()
	# print("uSign public key base64: " + str(str_base64_pub_check))
	uCheck = LibuECC()
	uCheck.writePublicKeyBase64(str_base64_pub_check)

	#region check sign
	str_correct = "this message will be signed"
	str_wrong = "this message is not signed"
	str_signed_data = uSign.Sign(str_correct)
	print("check correct result: " + str(uCheck.Verify(str_correct, str_signed_data)))
	print("check wrong result: " + str(uCheck.Verify(str_wrong, str_signed_data)))
	#endregion

	#region check ecdhe
	uRemote = LibuECC()
	uRemote.generateKeyPair()
	str_base64_pub_remote = uRemote.readPublicKeyBase64()
	# print("uRemote public key base64: " + str(str_base64_pub_remote))
	
	remoteShared = uRemote.SharedSecret(str_base64_pub_check)
	localShared = uSign.SharedSecret(str_base64_pub_remote)
	if remoteShared == localShared:
		print("ECDHE check passed")
	else:
		print("ECDHE failed")
	#endregion

	#region encryped and decrypt
	encrypted = uSign.Encrypt(str_correct);
	print("encrypted: " + str(encrypted));
	decrypted = uRemote.Decrype(encrypted);
	print("decryped: " + decrypted);
	#endregion

